import os
import sys
import time
import traceback
import json
from pathlib import Path, PurePath
from typing import Optional, Dict, Any
from dataclasses import dataclass, asdict

from dotenv import load_dotenv
from openai import OpenAI, OpenAIError
from openai.types import Video


from utils import sr_s3
from utils.sr_utils import get_video_info

# 配置与初始化
load_dotenv(verbose=True, override=True)


# 获取项目根目录

PROJECT_ROOT = Path(__file__).parent.parent.parent.parent.absolute()
print("Project root:", PROJECT_ROOT)

# 环境变量与客户端初始化

# ============================================================
# 配置与常量
# ============================================================

DEFAULT_MODEL = "sora-2"
DEFAULT_RESOLUTIONS = {
    "sora-2": ["1280x720", "720x1280"],
    "sora-2-pro": ["1280x720", "720x1280", "1024x1792", "1792x1024"]
}
DEFAULT_DURATION = "8"
DEFAULT_DOWNLOAD_DIR = "download/videos"
DEFAULT_RESOLUTION = {
    "9:16": "720x1280",
    "16:9": "1280x720"
}


# ============================================================
# 数据结构定义
# ============================================================

@dataclass
class VideoResult:
    status: int
    width: int = 0
    height: int = 0
    fps: int = 0
    url: Optional[str] = None
    local_path: Optional[str] = None
    msg: str = ""
    code: str = ""

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)
    
    def to_json(self) -> str:
        return json.dumps(self.to_dict(), ensure_ascii=False)


# ============================================================
# 上传模块
# ============================================================

def upload_video_to_s3(video_path: str, s3_storage: str, storage_service: str, bucket_name: str) -> tuple[Optional[str], Optional[str]]:
    """上传视频到S3 (R2) 并返回路径"""
    if not video_path or not os.path.exists(video_path):
        return None, "视频文件不存在"

    try:
        s3_key = sr_s3.upload_login_or_visitor(
            video_path, s3_storage, serviceName=storage_service, bucketName=bucket_name
        )
        return s3_key, None
    except Exception as e:
        return None, f"上传R2失败: {str(e)}"


# ============================================================
# 视频生成器类
# ============================================================

class Sora2VideoGenerator:
    """Sora-2 视频生成API封装类"""

    def __init__(self, api_key: Optional[str] = None):
        self.OPENAI_API_KEY = api_key or os.getenv("OPENAI_API_KEY", "")
        self.client = OpenAI(api_key=self.OPENAI_API_KEY)

    def create_video_job(self, prompt: str,
                         image: Optional[Path] = None,
                         model: str = DEFAULT_MODEL,
                         size: str = "720x1280",
                         seconds: str = DEFAULT_DURATION) -> Video:
        """提交视频生成任务"""
        try:
            payload = {"model": model, "prompt": prompt, "size": size, "seconds": seconds}
            if image and image.exists():
                # 如果提供了图片，先调整图片尺寸为目标尺寸，保持宽高比，不足部分填充黑色
                from PIL import Image

                # 解析目标尺寸
                target_width, target_height = map(int, size.split('x'))

                resized_image_path = image.with_name(f"resized_{image.name}")
                with Image.open(image) as img:
                    # 将图片转换为RGB模式（去除alpha通道）
                    if img.mode != 'RGB':
                        img = img.convert('RGB')

                    # 计算调整比例，保持宽高比，使图片最大边契合目标尺寸
                    img_width, img_height = img.size
                    
                    # 根据目标尺寸和原图比例选择调整策略
                    if target_width >= target_height:  # 横屏或方屏
                        if img_width >= img_height:  # 原图也是横屏或方屏
                            scale = target_width / img_width
                        else:  # 原图为竖屏
                            scale = target_height / img_height
                    else:  # 竖屏
                        if img_width >= img_height:  # 原图为横屏
                            scale = target_width / img_width
                        else:  # 原图也是竖屏
                            scale = target_height / img_height

                    # 计算缩放后的尺寸
                    new_width = int(img_width * scale)
                    new_height = int(img_height * scale)

                    print(f"调整比例: {scale:.2f}, 调整后的尺寸: {new_width}x{new_height}")
                    # 调整图片
                    resized_img = img.resize((new_width, new_height), Image.LANCZOS)

                    # 创建黑色背景图片
                    final_img = Image.new('RGB', (target_width, target_height), (0, 0, 0))

                    # 计算居中位置
                    x = (target_width - new_width) // 2
                    y = (target_height - new_height) // 2

                    # 将缩放后的图片粘贴到黑色背景上
                    final_img.paste(resized_img, (x, y))

                    # 保存处理后的图片
                    final_img.save(resized_image_path)

                # 使用调整后的图片
                with open(resized_image_path, "rb") as img_file:
                    payload["input_reference"] = img_file
                    result = self.client.videos.create(**payload)

                # 删除临时调整尺寸的图片
                try:
                    os.remove(resized_image_path)
                except:
                    pass

                return result
            return self.client.videos.create(**payload)
        except OpenAIError as e:
            print("Error when creating video:", e)
            raise

    def get_video_job(self, video_id: str) -> Video:
        try:
            return self.client.videos.retrieve(video_id)
        except OpenAIError as e:
            print("Error when retrieving video job:", e)
            raise

    def download_video_content(self, video_id: str, out_path: str) -> str:
        """下载视频到本地"""
        try:
            # 确保目录存在，基于项目根目录
            # 移除路径开头的 "./" 或 "/" 等符号
            clean_path = out_path.lstrip('./').lstrip('/')
            full_out_path = PROJECT_ROOT / clean_path
            print("Project root:", PROJECT_ROOT)
            print("Clean path:", clean_path)
            print("Full out path:", full_out_path)
            os.makedirs(os.path.dirname(full_out_path), exist_ok=True)
            content = self.client.videos.download_content(video_id=video_id, variant="video")
            print("Downloading video content to:", full_out_path.resolve())
            # 确保full_out_path是文件路径而不是目录路径
            content.write_to_file(full_out_path)
            # 返回规范化的相对路径
            return full_out_path
        except OpenAIError as e:
            print("Error downloading video content:", e)
            raise

    def process_video_job(self, prompt: str,
                          image: Optional[Path] = None,
                          model: str = DEFAULT_MODEL,
                          size: str = "720x1280",
                          seconds: str = DEFAULT_DURATION,
                          download_path: str = DEFAULT_DOWNLOAD_DIR) -> str:
        """提交任务并轮询进度"""
        job = self.create_video_job(prompt, image, model, size, seconds)
        job_id = getattr(job, "id", None)
        # job_id = "video_68f0f2c3a14c81988d0e1ea849a69bb106b79084d2cbd7df"
        print(f"Video generation started: {job_id}")

        start_time = time.time()
        while True:
            info = self.get_video_job(job_id)
            status = getattr(info, "status", "unknown")
            # progress = getattr(info, "progress", 0)
            # sys.stdout.write(f"\r[{status}] {progress:.1f}%")
            # sys.stdout.flush()

            # 检查是否超时（20分钟）
            if time.time() - start_time > 20 * 60:
                return VideoResult(status=0, msg="Video generation timeout", code="6009").to_json()

            if status in ("failed", "completed"):
                break
            time.sleep(10)
        # sys.stdout.write("\n")

        if status == "failed":
            error = getattr(info, "error", None)
            if error and getattr(error, "code", "") == "moderation_blocked":
                msg = getattr(error, "message", "Your request was blocked by our moderation system.")
                return VideoResult(status=0, msg=msg, code="6008").to_json()

            msg = getattr(error, "message", "Video generation failed")
            return VideoResult(status=0, msg=msg, code="6001").to_json()

        # 构建相对于项目根目录的完整路径
        filename = f"{download_path}/video_task_{job_id}.mp4"
        out_file = self.download_video_content(job_id, filename)
        print(f"Video downloaded to2: {out_file}")
        video_info = get_video_info(out_file)

        return VideoResult(
            status=1,
            width=video_info.get("width", 0),
            height=video_info.get("height", 0),
            fps=video_info.get("fps", 0),
            local_path=str(out_file),
            msg="success"
        ).to_json()


# ============================================================
# 主处理函数
# ============================================================

def generator_openai(body: Dict[str, Any], local_time=None) -> Dict[str, Any]:
    """统一处理视频生成任务"""
    task_id = body.get("taskId", "unknown")
    try:
        # 参数校验
        if not body.get("prompt"):
            return VideoResult(status=0, msg="缺少参数: prompt", code="6001").to_dict()
        if not body.get("moduleName"):
            return VideoResult(status=0, msg="缺少参数: moduleName", code="6001").to_dict()

        if not body.get("ratio") or not body.get("ratio").strip():
            body["ratio"] = "16:9"
        prompt = body["prompt"]
        moduleName = body["moduleName"]
        ratio = body["ratio"]
        duration = str(body.get("duration", DEFAULT_DURATION))
        print("图片地址1：", body.get("first_locaPath"))
        first_image_path = body.get("first_locaPath")
        # 检查first_image是否为空字符串或None
        if not first_image_path or not first_image_path.strip():
            first_image = None
            print("图片参数为空或仅包含空格，设置为None")
        else:
            # 处理可能包含查询参数的URL，只保留路径部分
            if '?' in first_image_path:
                first_image_path = first_image_path.split('?')[0]
                print(f"移除查询参数后的图片路径: {first_image_path}")

            first_image = Path(first_image_path)
            if not first_image.exists():
                print(f"图片文件不存在: {first_image}")
                first_image = None
            else:
                print(f"图片文件存在: {first_image}")
        print("图片地址2：", first_image)

        # 根据模型类型选择合适的分辨率
        if moduleName in DEFAULT_RESOLUTIONS:
            # 检查请求的分辨率是否支持该模型
            size = DEFAULT_RESOLUTION[ratio]
            if size not in DEFAULT_RESOLUTIONS[moduleName]:
                # 如果不支持，则使用该模型的第一个支持的分辨率
                size = DEFAULT_RESOLUTIONS[moduleName][0]
        else:
            # 如果模型不在配置中，使用默认分辨率
            size = DEFAULT_RESOLUTION[ratio]
            
        # 视频生成
        generator = Sora2VideoGenerator()
        result_json = generator.process_video_job(
            prompt=prompt,
            model=moduleName,
            size=size,
            image=first_image,
            seconds=duration
        )
        print(f"任务: {type(result_json)}")
        result = json.loads(result_json)

        if not result.get("local_path"):
            return VideoResult(status=0, msg=result.get("msg", "生成失败"), code=result.get("code", "6001")).to_dict()

        # 上传
        print(f"开始上传视频到S3: {result['local_path']}")
        s3_key, err = upload_video_to_s3(
            result["local_path"],
            body.get("storedPrefix"),
            body.get("storageService", "s3"),
            body.get("bucketName")
        )
        if not s3_key:
            return VideoResult(status=0, msg=f"视频上传R2失败: {err}").to_dict()

        return VideoResult(
            status=1,
            width=result.get("width", 0),
            height=result.get("height", 0),
            fps=result.get("fps", 0),
            url=s3_key,
            local_path=result["local_path"],
            msg="success"
        ).to_dict()

    except Exception as e:
        print(f"任务 {task_id} 异常: {traceback.format_exc()}")
        return VideoResult(status=0, msg=f"任务处理异常: {str(e)}", code="6001").to_dict()


# ============================================================
# 调试入口
# ============================================================

if __name__ == "__main__":
    task_params = {
        "taskId": "17601299299252jW0P1564wHMH63h",
        "prompt": "As if the baby on the woman's lap was coming down and standing on the ground",
        "first_image": "video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1760129915983.jpeg",
        "duration": "5",
        "storageService": "r2",
        "bucketName": "picwand",
        'ratio': '16:9',
        'moduleName': 'sora-2',
        "storedPrefix": "video/image_to_video",
    }
    result = generator_openai(task_params)
    print(result)